//! Types for manipulating ararys.

/// Describes a "square" bounded multidimensional array shape.
#[derive(Clone, Debug)]
pub struct ArrayShape {
    /// The dimensions of each chunk.
    ///
    /// The innermost arrays are first, the outermost last.  Think of this as
    /// wrapping an array of some topology in another array which adds another
    /// dimension.  Those outer arrays are indexed "first".
    dim_sizes: Vec<u64>,
}

impl ArrayShape {
    /// Creates new shape with a single unit value.
    pub fn new_unit() -> Self {
        Self {
            dim_sizes: Vec::new(),
        }
    }

    /// Takes this array dimensions and constructs a new array.
    pub fn in_array_of_size(mut self, size: u64) -> Self {
        self.dim_sizes.push(size);
        self
    }

    /// Creates a array shape of a flat size, which is the most common shape.
    pub fn flat(size: u64) -> Self {
        Self::new_unit().in_array_of_size(size)
    }

    /// Returns if this is a unit shape.
    pub fn is_unit(&self) -> bool {
        self.dim_sizes.is_empty()
    }

    /// Returns the number of dimensions the array is.
    pub fn dimensions(&self) -> usize {
        self.dim_sizes.len() as usize
    }

    /// Returns the total number of unit values in the array structure.
    pub fn total_elements(&self) -> u64 {
        self.dim_sizes.iter().fold(1, |a, v| a * v)
    }

    /// Gets the size of the dimension of some depth.
    ///
    /// Checking for a depth that is the total number of dimensions will return
    /// 0, as if a scalar is a "0 dimensional" array.
    ///
    /// # Panics
    ///
    /// If querying a depth too deep.
    pub fn depth_size(&self, depth: usize) -> u64 {
        if depth > self.dimensions() {
            panic!("accessor: dim depth too deep");
        }

        if depth == self.dimensions() {
            return 0;
        }

        self.dim_sizes[self.dimensions() - depth - 1]
    }

    /// Returns the stride size for queries of the given depth.
    ///
    /// A depth of 0 is just the whole array, which is useful for sequences
    /// which would the variable length arrays around this one.
    ///
    /// # Panics
    ///
    /// If provided with a depth deeper than the number of dimensions we have.
    pub fn depth_stride_size(&self, depth: usize) -> u64 {
        if depth > self.dimensions() {
            panic!("accessor: stride depth too deep");
        }

        let to_take = self.dimensions() - depth;
        self.dim_sizes.iter().take(to_take).fold(1, |a, v| a * v)
    }

    /// Checks to see if an index is in bounds for this shape.
    ///
    /// If it is, returns `None`, if it's not, returns the index of the
    /// shallowest failed query.
    ///
    /// # Panics
    ///
    /// If provided with a coordinate with more dimensions than we have in this
    /// shape.
    pub fn check_index_in_bounds(&self, coord: &[u64]) -> Option<usize> {
        if coord.len() > self.dimensions() {
            panic!("accessor: coordinate query index too deep");
        }

        for (i, c) in coord.iter().enumerate() {
            let dim_idx = self.dimensions() - i - 1;
            let dim_size = self.dim_sizes[dim_idx];
            if *c >= dim_size {
                return Some(i);
            }
        }

        None
    }

    /// Iterates over the dimension lengths, largest stride first.
    pub fn dims_iter(&self) -> impl Iterator<Item = u64> + '_ {
        self.dim_sizes.iter().rev().copied()
    }
}

#[cfg(test)]
mod tests {
    use super::ArrayShape;

    #[test]
    fn test_array_shape_size() {
        let shape = ArrayShape::new_unit();
        assert_eq!(shape.total_elements(), 1);

        let shape = shape.in_array_of_size(3);
        assert_eq!(shape.total_elements(), 3);

        let shape = shape.in_array_of_size(5);
        assert_eq!(shape.total_elements(), 15);

        let shape = shape.in_array_of_size(7);
        assert_eq!(shape.total_elements(), 105);
        assert_eq!(shape.dimensions(), 3);
    }

    #[test]
    fn test_array_shape_stride() {
        let shape = ArrayShape::new_unit();
        assert_eq!(shape.depth_stride_size(0), 1);

        let shape = shape.in_array_of_size(3);
        assert_eq!(shape.depth_stride_size(0), 3);
        assert_eq!(shape.depth_stride_size(1), 1);

        let shape = shape.in_array_of_size(5);
        assert_eq!(shape.depth_stride_size(0), 15);
        assert_eq!(shape.depth_stride_size(1), 3);
        assert_eq!(shape.depth_stride_size(2), 1);

        let shape = shape.in_array_of_size(7);
        assert_eq!(shape.depth_stride_size(0), 105);
        assert_eq!(shape.depth_stride_size(1), 15);
        assert_eq!(shape.depth_stride_size(2), 3);
        assert_eq!(shape.depth_stride_size(3), 1);
    }
}
